#!/usr/bin/python3
# coding=utf-8

# 批量运行 pcoc
pcoc_text_file = 'pt'

work_file = 'tmpsh'

cmd_pcoc_docker = '''docker run -e LOCAL_USER_ID=`id -u $USER` --rm -v $PWD:$PWD -e CWD=$PWD carinerey/pcoc'''

import shutil

import os
import threading

import sys
import argparse


parser = argparse.ArgumentParser(
    description=''' 批量运行pcoc
    用法:
    conver03_run_pcoc.py  -i OGS_for_second_step -t 64
    由大天才于2021年11月12日创建于浙江农业大学''')





parser.add_argument('-i',
                help='必须给定，工作路径')

parser.add_argument('-t',
                help='并行核数 默认为8')


args = parser.parse_args()

try:
	shutil.rmtree(work_file)
except:
	pass
try:
	os.mkdir(work_file)
except:
	pass


if not args.i:
    parser.print_help()
    sys.exit()


infile = args.i

#infile  = 'test'

print('注意 这一步需要在root下使用docker')

if not args.t:
	thread = 8
else:
	thread = int(args.t)

#thread = 5





thread_dic = {i:[] for i in range(0,thread)}

n = 0
for i in os.listdir(infile):
	if os.path.isdir(infile+'/'+i):
		thread_dic[n%thread].append(i)
		n+=1

cwd = os.getcwd()

for i in thread_dic:

	outfile = open(work_file+'/'+str(i)+'.sh','w')

	for j in thread_dic[i]:
		outfile.write('cd '+cwd+'/'+infile+'/'+j+'\n')
		# 读取pt text
		with open(infile+'/'+j+'/'+pcoc_text_file) as fila:
			pt_text = next(fila).strip()
			if pt_text!='':

				outfile.write(cmd_pcoc_docker + ''' pcoc_det.py -t spe_tree -aa pep_cut_ali.fa -o output_pcoc_det -m %s --plot_complete_ali --plot\n''' % pt_text)
			


		outfile.write('cd '+cwd +'\n')
	outfile.close()



def run_func(file_path):


	a = os.system('sh '+file_path)


	return	None

n = 0
thread_pool = []
for i in thread_dic:

	file_path = work_file+'/'+str(i)+'.sh'

	n += 1

	th = threading.Thread(target=run_func, args=[file_path], name='th_'+str(n))
	thread_pool.append(th)


for i in range(thread):
	if len(thread_pool)>i:
		thread_pool[i].start()
